import os
import random
import skvideo.io
import cv2

random.seed(123)

VER = 'v1.0'
BASE_DIR, FILE_NAME = os.path.split(__file__)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)

# VIDEO_SOURCE = r"C:\DL_DATA\dl2_DataSets\traffic_counter\traffic_1min.mp4"
path = '../../../../large_data/video/traffic/traffic_1min.mp4'
VIDEO_SOURCE = os.path.join(BASE_DIR, path)
# IMAGE_DIR = r"C:\DL_DATA\dl2_DataSets\traffic_counter\out1"
IMAGE_DIR = os.path.join(SAVE_DIR, 'out1')
# out_path = './out'
out_path = os.path.join(SAVE_DIR, 'out')

SHAPE = (720, 1280)  # HxW


# 背景扣除算法
def train_bg_subtractor(inst, cap, num=500):
    print('Training BG Subtractor...')
    i = 0
    for frame in cap:
        inst.apply(frame, None, 0.001)
        i += 1
        if i >= num:
            return cap
    print(f'Warning: not enough frames to read: asked for {num}, only {i}.')
    return cap


# 滤波
def filter_mask(img, a=None):
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))
    # Fill any small holes
    closing = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
    # Remove noise
    opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, kernel)
    # Dilate to merge adjacent blobs
    dilation = cv2.dilate(opening, kernel, iterations=2)

    return dilation


bg_subtractor = cv2.createBackgroundSubtractorMOG2(history=500, detectShadows=True)
cap = skvideo.io.vreader(VIDEO_SOURCE)
train_bg_subtractor(bg_subtractor, cap, num=500)

frame_number = -1
for frame in cap:
    frame_number += 1
    fg_mask = bg_subtractor.apply(frame, None, 0.001)
    fg_mask[fg_mask < 240] = 0
    fg_mask = filter_mask(fg_mask, frame_number)
    cv2.imwrite(out_path + '/' + "fg_mask_%04d.png" % frame_number, fg_mask)

img = cv2.imread(out_path + '/' + 'fg_mask_0000.png')
print(img.shape)

fps = 48
size = (img.shape[1], img.shape[0])

fourcc = cv2.VideoWriter_fourcc(*'XVID')
videoWriter = cv2.VideoWriter('output_fg.avi', fourcc, fps, size)

for i in os.listdir(out_path):
    im_name = os.path.join(out_path, i)
    frame = cv2.imread(im_name)
    videoWriter.write(frame)

videoWriter.release()
print('finish')
